{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pygame 1.9.6\n",
      "Hello from the pygame community. https://www.pygame.org/contribute.html\n",
      "Start collecting images...\n",
      "Done\n"
     ]
    },
    {
     "ename": "error",
     "evalue": "unpack requires a buffer of 4 bytes",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31merror\u001b[0m                                     Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-1-28aad307b749>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m    188\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    189\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0m__name__\u001b[0m \u001b[1;33m==\u001b[0m \u001b[1;34m'__main__'\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 190\u001b[1;33m     \u001b[0mCollectTrainingData\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32m<ipython-input-1-28aad307b749>\u001b[0m in \u001b[0;36m__init__\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m     37\u001b[0m         \u001b[0mpygame\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0minit\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     38\u001b[0m         \u001b[0mpygame\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdisplay\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mset_mode\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m400\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m300\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 39\u001b[1;33m         \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcollect_image\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     40\u001b[0m     \u001b[1;31m#收集图像\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     41\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mcollect_image\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-1-28aad307b749>\u001b[0m in \u001b[0;36mcollect_image\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m     57\u001b[0m             \u001b[1;32mwhile\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msend_inst\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     58\u001b[0m                     \u001b[1;31m# 获取数据帧\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 59\u001b[1;33m                     \u001b[0mimage_len\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mstruct\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0munpack\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'<L'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mconnection\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mread\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mstruct\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcalcsize\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'<L'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     60\u001b[0m                     \u001b[1;32mif\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[0mimage_len\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     61\u001b[0m                         \u001b[1;32mbreak\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31merror\u001b[0m: unpack requires a buffer of 4 bytes"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import os\n",
    "import sys\n",
    "import socket\n",
    "import struct\n",
    "import serial\n",
    "import numpy as np\n",
    "import pygame\n",
    "import cv2\n",
    "#from serial import Serial\n",
    "\n",
    "\n",
    "class CollectTrainingData(object):\n",
    "    \n",
    "    def __init__(self):\n",
    "\n",
    "        HOST = '0.0.0.0'\n",
    "        PORT = 10000\n",
    "        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)\n",
    "        # 启动socket，设置监听端口为10000，接受所有ip的连接\n",
    "        self.sock.bind((HOST, PORT))\n",
    "        self.sock.listen(1)\n",
    "        # 接受一个客户端连接\n",
    "        self.connection = self.sock.accept()[0].makefile('rb')       \n",
    "        \n",
    "        \n",
    "        #连接到串口\n",
    "        self.ser = serial.Serial('COM5', 9600, timeout=1)\n",
    "        self.send_inst = True\n",
    "       \n",
    "        # 创建标签\n",
    "        self.k = np.zeros((4, 4), 'float')\n",
    "        for i in range(4):\n",
    "            self.k[i, i] = 1\n",
    "        self.temp_label = np.zeros((1, 4), 'float')\n",
    "\n",
    "        pygame.init()\n",
    "        pygame.display.set_mode((400, 300))\n",
    "        self.collect_image()\n",
    "    #收集图像\n",
    "    def collect_image(self):\n",
    "\n",
    "        saved_frame = 0   #已经保存的帧\n",
    "        total_frame = 0   #所有的帧\n",
    "\n",
    "        # 从摄像头收集图像\n",
    "        print('Start collecting images...')\n",
    "        e1 = cv2.getTickCount()\n",
    "        image_array = np.zeros((1, 38400))      #图像的大小是 1 * 38400   \n",
    "        label_array = np.zeros((1, 4), 'float') #图像标签的大小是 1 * 4 \n",
    "\n",
    "        #  逐帧获取图像\n",
    "        try:\n",
    "            frame = 1\n",
    "            \n",
    "\n",
    "            while self.send_inst:\n",
    "                    # 获取数据帧\n",
    "                    image_len = struct.unpack('<L', self.connection.read(struct.calcsize('<L')))[0]\n",
    "                    if not image_len:\n",
    "                        break\n",
    "\n",
    "                    recv_bytes = b''\n",
    "                    recv_bytes += self.connection.read(image_len)\n",
    "                    image = cv2.imdecode(np.frombuffer(recv_bytes, dtype=np.uint8), cv2.IMREAD_GRAYSCALE)\n",
    "                    # 选择图片的下半部分\n",
    "                    roi = image[120:240, :]\n",
    "                    \n",
    "                    # 保存整张图片\n",
    "                    cv2.imwrite('training_images/frame{:>05}.jpg'.format(frame), image)\n",
    "                    \n",
    "                    #展示整张图片\n",
    "                    cv2.imshow('image', image)\n",
    "                    \n",
    "                    # 下半部分的图片转换成一维\n",
    "                    temp_array = roi.reshape(1, 38400).astype(np.float32)\n",
    "                    \n",
    "                    frame += 1\n",
    "                    total_frame += 1\n",
    "                    #如果输入是指定小车运动的6个操作之一，\n",
    "                    #那么这个frame 就保存起来，saved_frame += 1，\n",
    "                    #而不管输入的操作是什么，`total_frame += 1`\n",
    "\n",
    "                    # 获取用户的输入并执行\n",
    "                    for event in pygame.event.get():\n",
    "                        if event.type == pygame.KEYDOWN:\n",
    "                            key_input = pygame.key.get_pressed()\n",
    "\n",
    "                            # 复合操作\n",
    "                            if key_input[pygame.K_o]and key_input[pygame.K_p]:\n",
    "                                print(\"Forward Right\")\n",
    "                                image_array = np.vstack((image_array, temp_array))\n",
    "                                #np.vstack:按垂直方向（行顺序）堆叠数组构成一个新的数组\n",
    "                                label_array = np.vstack((label_array, self.k[1]))\n",
    "                                #拼成一个整图片，上半部分是空白的， 下半部分是从摄像头中获取的，标签也类似\n",
    "                                saved_frame += 1\n",
    "                                self.ser.write(b'5')\n",
    "\n",
    "                            elif key_input[pygame.K_i]and key_input[pygame.K_u]:\n",
    "                                print(\"Forward Left\")\n",
    "                                image_array = np.vstack((image_array, temp_array))\n",
    "                                label_array = np.vstack((label_array, self.k[0]))\n",
    "                                saved_frame += 1\n",
    "                                self.ser.write(b'6')\n",
    "\n",
    "                            elif key_input[pygame.K_DOWN] and key_input[pygame.K_RIGHT]:\n",
    "                                print(\"Reverse Right\")\n",
    "                                self.ser.write(b'7')\n",
    "                            \n",
    "                            elif key_input[pygame.K_DOWN] and key_input[pygame.K_LEFT]:\n",
    "                                print(\"Reverse Left\")\n",
    "                                self.ser.write(b'8')\n",
    "\n",
    "                            # 单个操作\n",
    "                            elif key_input[pygame.K_UP]:\n",
    "                                print(\"Forward\")\n",
    "                                saved_frame += 1\n",
    "                                image_array = np.vstack((image_array, temp_array))\n",
    "                                label_array = np.vstack((label_array, self.k[2]))\n",
    "                                self.ser.write(b'51')\n",
    "\n",
    "                            elif key_input[pygame.K_DOWN]:\n",
    "                                print(\"Reverse\")\n",
    "                                saved_frame += 1\n",
    "                                image_array = np.vstack((image_array, temp_array))\n",
    "                                label_array = np.vstack((label_array, self.k[3]))\n",
    "                                self.ser.write(b'52')\n",
    "                            \n",
    "                            elif key_input[pygame.K_RIGHT]:\n",
    "                                print(\"Right\")\n",
    "                                image_array = np.vstack((image_array, temp_array))\n",
    "                                label_array = np.vstack((label_array, self.k[1]))\n",
    "                                saved_frame += 1\n",
    "                                self.ser.write(b'54')\n",
    "\n",
    "                            elif key_input[pygame.K_LEFT]:\n",
    "                                print(\"Left\")\n",
    "                                image_array = np.vstack((image_array, temp_array))\n",
    "                                label_array = np.vstack((label_array, self.k[0]))\n",
    "                                saved_frame += 1\n",
    "                                self.ser.write(b'53')\n",
    "\n",
    "                            elif key_input[pygame.K_x] or key_input[pygame.K_q]:\n",
    "                                print('exit')\n",
    "                                self.send_inst = False\n",
    "                                self.ser.write(b'0')\n",
    "                                break\n",
    "                                    \n",
    "                        elif event.type == pygame.KEYUP:\n",
    "                            self.ser.write(b'0')\n",
    "\n",
    "            # save training images and labels\n",
    "            train = image_array[1:, :]\n",
    "            train_labels = label_array[1:, :]\n",
    "\n",
    "            # 把train的图片和label转换为numpy，保存\n",
    "            file_name = str(int(time.time()))\n",
    "            directory = \"training_data\"\n",
    "            if not os.path.exists(directory):\n",
    "                os.makedirs(directory)\n",
    "            try:    \n",
    "                np.savez(directory + '/' + file_name + '.npz', train=train, train_labels=train_labels)\n",
    "            #如果你想将多个数组保存到一个文件中的话，可以使用numpy.savez函数。\n",
    "            #savez函数的第一个参数是文件名，其后的参数都是需要保存的数组，\n",
    "            #也可以使用关键字参数为数组起一个名字，\n",
    "            #非关键字参数传递的数组会自动起名为arr_0, arr_1, …。\n",
    "            #savez函数输出的是一个压缩文件(扩展名为npz)，\n",
    "            #其中每个文件都是一个save函数保存的npy文件，文件名对应于数组名。\n",
    "            #load函数自动识别npz文件，并且返回一个类似于字典的对象，\n",
    "            #可以通过数组名作为关键字获取数组的内容：\n",
    "\n",
    "            except IOError as e:\n",
    "                print(e)\n",
    "\n",
    "            e2 = cv2.getTickCount()\n",
    "            # 计算服务持续的时间\n",
    "            time0 = (e2 - e1) / cv2.getTickFrequency()\n",
    "            print('Streaming duration:', time0)\n",
    "\n",
    "            print((train.shape))\n",
    "            print((train_labels.shape))\n",
    "            print('Total frame:', total_frame)\n",
    "            print('Saved frame:', saved_frame)\n",
    "            print('Dropped frame', total_frame - saved_frame)\n",
    "\n",
    "        finally:\n",
    "            print('Done')\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    CollectTrainingData()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
